LlamaIndexを完全に理解するチュートリアルその２：テキスト分割のカスタマイズ

第２回はテキスト分割を脱ブラックボックス化し、カスタマイズするぞ！

#LlamaIndex

#ChatGPT

#OpenAI

#機械学習

#Python

nokomoro3

2023.05.27

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

こんちには。

データアナリティクス事業本部インテグレーション部機械学習チームの中村です。

「LlamaIndexを完全に理解するチュートリアルその２」では、テキスト分割のカスタマイズを取り扱います。

本記事で使用する用語は以下のその１で説明していますので、そちらも参照ください。

LlamaIndexを完全に理解するチュートリアル
その１：処理の概念や流れを理解する基礎編（v0.7.9対応）
その２：テキスト分割のカスタマイズ

・本記事の内容はその１のv0.7.9版の記事を投稿後、v0.7.9で動作するように修正しています

本記事の内容

LlamaIndexは与えられたテキストをインデックス化しますが、LLMとのやり取りではテキストの長さ（トークン数）の上限という制約があります。

具体的に、ChatGPTの裏側で使用されるgpt-3.5-turboなどは、トークン数が4096個が上限となり、日本語の文字数でもおおよそ同じ数が上限となります。

この制約は、「入力（プロンプト）」と「応答」双方の合計値に対するものですので、文脈を与えてチャットに応答される場合、与えられる文脈には制限が出てきます。

そのため、LlamaIndexでは、テキストをある単位でチャンク分割する処理を行います。

通常デフォルトで使用する場合は、チャンク分割が暗黙的に行われていますので、今回はこのカスタマイズ方法を見ていきます。

環境準備

その１と同様の方法で準備します。

使用したバージョン情報は以下となります。

Python : 3.10.11
langchain : 0.0.234
llama-index : 0.7.9
openai : 0.27.8

サンプルコード

ベースのサンプルは以下とします。こちらを実行しておきます。

from llama_index import SimpleDirectoryReader
from llama_index import ListIndex

documents = SimpleDirectoryReader(input_dir="./data").load_data()

list_index = ListIndex.from_documents(documents)

query_engine = list_index.as_query_engine()

ノード分割の状況を可視化

上記を実行後に、以下でノードの分割状況を可視化できます。

for doc_id, node in list_index.storage_context.docstore.docs.items():
    node_dict = node.__dict__
    print(f'{doc_id=}, len={len(node_dict["text"])}, start={node_dict["start_char_idx"]}, end={node_dict["end_char_idx"]}')

doc_id='a8ecf0bb-e664-4ba3-a77d-ec9b7d591b3e', len=903, start=0, end=903
doc_id='b4e33e3f-5a82-4076-acb8-2af54add4c46', len=899, start=904, end=1803
doc_id='c9b3c44b-c429-4911-bdc7-4146f2dc0de3', len=942, start=1798, end=2740
doc_id='a280f12a-2c37-4ae7-93c9-42b675453422', len=245, start=2747, end=2992
doc_id='e0296a5b-a7b3-4885-8e44-94f7bc845fb9', len=735, start=0, end=735
doc_id='0688f16e-1627-416a-bba9-2dd526f09dad', len=795, start=736, end=1531
doc_id='b18bcc33-2f3f-4c66-9c2c-96be8c306e4a', len=832, start=1532, end=2364
doc_id='ac537ef0-a364-4866-99c2-5efcabc58a6b', len=616, start=2365, end=2981

こちらがデフォルトの動作です。

カスタマイズするためのコード修正

チャンク分割はServiceContextのNodeParserのTextSplitterが担っています。

実際にはTextSplitterはLangChainで定義されているクラスで、LlamaIndexはTextSplitterのサブクラスとしてTokenTextSplitterなどを定義しています。

カスタマイズするの場合は以下のように表にTextSplitterを出します。

from llama_index import SimpleDirectoryReader
from llama_index import Document
from llama_index import GPTListIndex
from llama_index import ServiceContext
from llama_index.node_parser import SimpleNodeParser
from llama_index.langchain_helpers.text_splitter import TokenTextSplitter
from llama_index.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE
import tiktoken

documents = SimpleDirectoryReader(input_dir="./data").load_data()

text_splitter = TokenTextSplitter(separator=" ", chunk_size=DEFAULT_CHUNK_SIZE
    , chunk_overlap=DEFAULT_CHUNK_OVERLAP
    , tokenizer=tiktoken.get_encoding("gpt2").encode)
node_parser = SimpleNodeParser(text_splitter=text_splitter)

service_context = ServiceContext.from_defaults(
    node_parser=node_parser
)

list_index = GPTListIndex.from_documents(documents
    , service_context=service_context)

最初はデフォルトの動作と同じパラメータでTextSplitterを作成しました。

再度可視化してみましょう。

for doc_id, node in list_index.storage_context.docstore.docs.items():
    node_dict = node.__dict__
    print(f'{doc_id=}, len={len(node_dict["text"])}, start={node_dict["start_char_idx"]}, end={node_dict["end_char_idx"]}')

doc_id='a8ecf0bb-e664-4ba3-a77d-ec9b7d591b3e', len=903, start=0, end=903
doc_id='b4e33e3f-5a82-4076-acb8-2af54add4c46', len=899, start=904, end=1803
doc_id='c9b3c44b-c429-4911-bdc7-4146f2dc0de3', len=942, start=1798, end=2740
doc_id='a280f12a-2c37-4ae7-93c9-42b675453422', len=245, start=2747, end=2992
doc_id='e0296a5b-a7b3-4885-8e44-94f7bc845fb9', len=735, start=0, end=735
doc_id='0688f16e-1627-416a-bba9-2dd526f09dad', len=795, start=736, end=1531
doc_id='b18bcc33-2f3f-4c66-9c2c-96be8c306e4a', len=832, start=1532, end=2364
doc_id='ac537ef0-a364-4866-99c2-5efcabc58a6b', len=616, start=2365, end=2981
doc_id='f1622414-e837-460a-8321-34eccfc61993', len=903, start=0, end=903
doc_id='69cee1eb-f457-45b1-9855-d6b259554dab', len=899, start=904, end=1803
doc_id='f1a4f4eb-9799-4389-bac8-e3aae306f945', len=942, start=1798, end=2740
doc_id='c9fe8605-f5e4-404d-9ac8-ca74892cb852', len=245, start=2747, end=2992
doc_id='e1c891b3-cdef-4227-8620-a6e317b89f64', len=735, start=0, end=735
doc_id='35b4cb29-e1bc-470b-b877-8f3026bf6440', len=795, start=736, end=1531
doc_id='256979b7-0bca-4048-b63b-32cf724a5faf', len=832, start=1532, end=2364
doc_id='10e81af7-e023-461d-bcc5-1db9a0b9db99', len=616, start=2365, end=2981

同じ結果となることが確認できました。

カスタマイズの詳細

テキスト分割は以下の部分に集約されます。

text_splitter = TokenTextSplitter(separator=" ", chunk_size=DEFAULT_CHUNK_SIZE
    , chunk_overlap=DEFAULT_CHUNK_OVERLAP
    , tokenizer=tiktoken.get_encoding("gpt2").encode)

各パラメータの意味は以下です。

separator : 区切り文字。ここの指定文字がチャンク分割の切れ目になるようチャンクを作ります。
chunk_size : チャンクサイズ。このチャンクサイズ以下となるようチャンクを作ります。
chunk_overlap : チャンクのオーバーラップ。後ろのチャンクに前のチャンクの末尾を付け加えることで、チャンク分割による文脈の断裂を緩和します。
tokenizer : tokenizerのencode処理。チャンクサイズをトークン数でカウントするために使用されます。

separatorについて

separatorはデフォルトは半角空白ですが、今回のサンプルデータでは句読点の後ろにたまたま半角空白が入っているため、句読点でうまく区切れているようです。

空白が無い場合で、句読点できちんと区切ってほしい場合は句読点を設定する必要がありそうです。

またseparatorはstr型に対するsplitと同じ動作で処理しますので、複数文字を設定すると、それをひと固まりの区切り文字として処理しますので、注意が必要です。

"hoge fuga. fuga hoge.".split(". ")

['hoge fuga', 'fuga hoge.']

ですので、ORで区切り文字を構成するには前処理で工夫するか、TokenTextSplitterでre.splitを使うようにサブクラスを再定義するかなどの必要がありそうですのでご注意ください。

tokenizerについて

チャンクサイズはトークン数単位でカウントするようですので、それにこのtokenizerが使用されます。

可視化をトークン数でやり直すと以下のようになります。

import tiktoken

for doc_id, node in list_index.storage_context.docstore.docs.items():
    node_dict = node.__dict__
    print(f'{doc_id=}, len={len(tiktoken.get_encoding("gpt2").encode(node_dict["text"]))}, start={node_dict["start_char_idx"]}, end={node_dict["end_char_idx"]}')

doc_id='f1ada414-fa64-4ca9-8c5f-7ac8b098c231', len=1013, start=0, end=903
doc_id='3b2bf3ae-c82d-4c1d-9f09-c3d5f77da1e1', len=1019, start=744, end=1651
doc_id='da959ec2-1932-49c7-b984-744306af0e54', len=964, start=1698, end=2567
doc_id='534ed4a4-b1a3-4237-b56b-3397c96f8707', len=807, start=2532, end=3272
doc_id='3bed0a84-4369-49f2-9d88-be757798a9fd', len=959, start=0, end=735
doc_id='b743ff61-09d1-4754-acb8-191045da82c3', len=997, start=621, end=1413
doc_id='64ef19cd-75fd-4cc4-88af-36b3be715c0d', len=1004, start=1374, end=2208
doc_id='204f050a-4410-4b84-a753-b10e5d65fc5a', len=991, start=2208, end=3032
doc_id='edb612e7-4c47-4c34-ae51-3b415eee2641', len=444, start=3069, end=3430

このようにチャンクサイズがDEFAULT_CHUNK_SIZE=1024以下のトークン数に収まっていることが分かります。

またデフォルトではgpt2のtokenizerが使用されている点も注意が必要です。

厳密には、tokenizerが後段のLLMやEmbeddingに合ったものを使用するのが理想的です。以下が代表的なものとなります。

text-davinci-003 : p50k_base
gpt-3.5-turbo : cl100k_base
gpt-4 : cl100k_base
text-embedding-ada-002 : cl100k_base

近年のモデルを使う場合、gpt2よりもcl100k_baseの方が適切かなという印象です。

詳細はtiktokenのコードを参照ください。

カスタマイズしてみた

以下のカスタマイズをした例を作ります。

tokenizerにcl100k_baseを使用
chunk_sizeを512に設定
chunk_overlapを100に設定
seperatorは。の句点と空白を設定

from llama_index import SimpleDirectoryReader
from llama_index import Document
from llama_index import GPTListIndex
from llama_index import ServiceContext
from llama_index.node_parser import SimpleNodeParser
from llama_index.langchain_helpers.text_splitter import TokenTextSplitter
from llama_index.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE
import tiktoken

documents = SimpleDirectoryReader(input_dir="./data").load_data()

text_splitter = TokenTextSplitter(separator="。 ", chunk_size=512
    , chunk_overlap=100
    , tokenizer=tiktoken.get_encoding("cl100k_base").encode)
node_parser = SimpleNodeParser(text_splitter=text_splitter)

service_context = ServiceContext.from_defaults(
    node_parser=node_parser
)

list_index = GPTListIndex.from_documents(documents
    , service_context=service_context)

念のため可視化してチェックします。

for doc_id, node in list_index.storage_context.docstore.docs.items():
    node_dict = node.__dict__
    print(f'{doc_id=}, len={len(node_dict["text"])}, start={node_dict["start_char_idx"]}, end={node_dict["end_char_idx"]}')

doc_id='5e5e6602-dbd3-4c73-9dd4-043d5b6151c5', len=554, start=0, end=554
doc_id='004d6b65-6d96-4655-9540-edd643e64858', len=570, start=458, end=1028
doc_id='d2efac46-5175-42df-8973-579c5585f69a', len=332, start=1049, end=1381
doc_id='036a7e5b-55e2-40b7-a74f-be3af525b1a7', len=467, start=1377, end=1844
doc_id='ce97efc6-9a8c-4220-8264-db1a906af1dc', len=480, start=1825, end=2305
doc_id='2a0cf7ae-ecbe-41c3-a6a4-97416c879fe1', len=554, start=2408, end=2962
doc_id='7b3cb425-bb3d-46fd-a1b4-decbbbfb2d34', len=414, start=2944, end=3358
doc_id='3bfb585f-7073-43e0-8bf0-e19e9a030705', len=492, start=0, end=492
doc_id='9de00bef-f127-42ee-aca6-56190de39c58', len=388, start=432, end=820
doc_id='c7367d75-9eb5-4a47-8cb4-68b35483a999', len=506, start=869, end=1375
doc_id='8ef855eb-faf8-462e-93f0-0e05bcdff09a', len=523, start=1307, end=1830
doc_id='1a2c124d-b64f-4009-a8ef-b89434e197e1', len=526, start=1822, end=2348
doc_id='cf265404-547d-4d6f-b944-07a9734c78cf', len=473, start=2397, end=2870
doc_id='8c2a40c0-6ba4-4039-92ae-d9ed65453cba', len=397, start=2900, end=3297

意図通りになっていることが分かりました。